import os
import shutil
import hashlib
import redis
from datetime import datetime
import re

# 对比文件签名，防重复
# compare_file_sing = False
compare_file_sing = True

# 使用文件名+尺寸的方式判断视频文件重复
use_filename_size_for_video_sign = False


# 文件整理器
# 根据redis配置判断文件重复
class FileOrganizer:
    def __init__(self, source_dir, target_dir, redis_host, redis_port, redis_db):
        self.source_dir = source_dir
        self.target_dir = target_dir
        self.file_hash_map = {
            "videos": {},
            "pictures": {},
            "documents": {},
            "archives": {}
        }
        self.redis_client = redis.Redis(host=redis_host, port=redis_port, db=redis_db)

    def organize_files(self):
        for root, dirs, files in os.walk(self.source_dir):
            # 忽略以.开头的文件和文件夹
            dirs[:] = [d for d in dirs if not d.startswith('.')]
            files = [f for f in files if not f.startswith('.')]

            for file in files:
                self.process_file(root, file)

    def process_file(self, root, file):
        file_path = os.path.join(root, file)
        file_name, file_extension = os.path.splitext(file)
        file_extension = file_extension.lower()

        # 预处理文件名
        if file_name.startswith('-'):
            file_name = file_name[1:]
        file_name = file_name.replace('-', '_')

        # 分类文件
        if file_extension in ['.mp4', '.mov', '.avi', '.rmi', '.rm', 'rmvb', '.mkv', 'wmv']:
            self.handle_file(file_path, file_name, file_extension, 'videos')
        elif file_extension in ['.jpg', '.jpeg', '.gif', '.png', '.cr2', '.cr3']:
            self.handle_file(file_path, file_name, file_extension, 'pictures')
        elif file_extension in ['.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pdf']:
            self.handle_file(file_path, file_name, file_extension, 'documents')
        elif file_extension in ['.zip', '.rar']:
            self.handle_file(file_path, file_name, file_extension, 'archives')

    def handle_file(self, file_path, file_name, file_extension, file_type):
        file_size = os.path.getsize(file_path)
        file_date = datetime.fromtimestamp(os.path.getmtime(file_path)).strftime('%Y%m')
        file_year = datetime.fromtimestamp(os.path.getmtime(file_path)).strftime('%Y')
        file_year_month = datetime.fromtimestamp(os.path.getmtime(file_path)).strftime('%Y%m')

        if file_type == 'videos':
            size_category = self.get_size_category(file_size)
            ext = file_extension[1:]
            target_path = os.path.join(self.target_dir, file_type, ext, f'{file_date}_{ext}_{size_category}',
                                       file_name + file_extension)
        elif file_type == 'pictures':
            ext = 'jpeg' if file_extension in ['.jpg', '.jpeg'] else file_extension[1:]
            target_path = os.path.join(self.target_dir, file_type, ext, file_year, file_year_month,
                                       file_name + file_extension)
        else:
            if file_extension in ['.doc', '.docs']:
                ext = 'doc'
            elif file_extension in ['.xls', '.xlsx']:
                ext = 'xls'
            elif file_extension in ['.ppt', '.pptx']:
                ext = 'ppt'
            else:
                ext = file_extension[1:]
            target_path = os.path.join(self.target_dir, file_type, ext, file_year_month,
                                       file_name + file_extension)

        self.save_file(file_path, target_path, file_type, file_extension)

    @staticmethod
    def get_size_category(size):
        if size < 200 * 1024 * 1024:
            return 'small'
        elif size < 800 * 1024 * 1024:
            return 'medium'
        else:
            return 'big'

    def save_file(self, source_path, target_path, file_type, file_extension):
        file_hash = self.get_file_hash(source_path, file_type, file_extension)

        if compare_file_sing and self.redis_client.sismember(f'{file_type}_hashes', file_hash):
            print(f"\tDuplicate file found: {source_path}")
            return

        if os.path.exists(target_path):
            target_path = self.resolve_duplicate(target_path)

        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        shutil.move(source_path, target_path)
        self.redis_client.sadd(f'{file_type}_hashes', file_hash)

        print(f"Moved {source_path} to {target_path}")

    @staticmethod
    def get_file_hash(file_path, file_type, file_extension):
        if use_filename_size_for_video_sign and file_type == 'videos':
            file_name = os.path.basename(file_path)
            file_size = os.path.getsize(file_path)
            return f'{file_size}_{file_name}_{file_extension}'
        else:
            hash_md5 = hashlib.md5()
            with open(file_path, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    hash_md5.update(chunk)
            return hash_md5.hexdigest()

    @staticmethod
    def resolve_duplicate(target_path):
        base, extension = os.path.splitext(target_path)
        for i in range(1, 100):
            new_path = f"{base}_{i:02d}{extension}"
            if not os.path.exists(new_path):
                return new_path
        raise Exception("Too many duplicate files")


class SpecifyFileOrganizer:
    def __init__(self, scan_dir, output_dir, classifications):
        """
        :param scan_dir: 要遍历的基目录
        :param output_dir: 文件移动后的输出目录
        :param classifications: 文件分类的字典，格式为 {"分类名": "正则表达式"}
        """
        self.scan_dir = scan_dir
        self.output_dir = output_dir
        self.classifications = classifications

    def organize_files(self):
        for root, dirs, files in os.walk(self.scan_dir):
            # 忽略以.开头的目录
            dirs[:] = [d for d in dirs if not d.startswith('.')]

            for file in files:
                if not file.startswith('.'):
                    self._process_file(root, file)

    def _process_file(self, root, file):
        file_path = os.path.join(root, file)
        for category, pattern in self.classifications.items():
            if re.match(pattern, file):
                file_date = datetime.fromtimestamp(os.path.getmtime(file_path)).strftime('%Y%m')
                self._move_file(file_path, file, category, file_date)
                break

    def _move_file(self, file_path, file_name, category, date_str):
        dest_dir = os.path.join(self.output_dir, category, date_str)
        os.makedirs(dest_dir, exist_ok=True)
        dest_path = os.path.join(dest_dir, file_name)

        print(f'Moving {file_path} to {dest_path}')
        shutil.move(file_path, dest_path)
        # shutil.copy(file_path, dest_path)
